package com.formula.distribution.limit.ratelimit.impl;

import com.formula.distribution.limit.enums.RateLimitConstantEnums;
import com.formula.distribution.limit.exception.RateLimitException;
import com.formula.distribution.limit.ratelimit.AbstractRateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import java.util.ArrayList;
import java.util.List;

/**
 * @author:luyanan
 * @email:luyanan0718@163.com
 * @date 2018/12/13 16:40
 * @introduce 令牌桶限流
 **/
@Slf4j
public class RedisRateLimiterTokenBucketImpl extends AbstractRateLimiter {


    @Autowired
    private RedisTemplate redisTemplate;


    private DefaultRedisScript redisScript;


    public RedisRateLimiterTokenBucketImpl(DefaultRedisScript redisScript) {
        this.redisScript = redisScript;
    }


    @Override
    public void rateLimit(String key, long limit, long refreshInterval, long tokenBucketStepNum, long tokenBucketTimeInterval, String message) {

        List<Object> keyList = new ArrayList<>();
        keyList.add(key);
        keyList.add(limit + RateLimitConstantEnums.HASH_TAG);
        keyList.add(tokenBucketStepNum + RateLimitConstantEnums.HASH_TAG);
        keyList.add(tokenBucketTimeInterval + RateLimitConstantEnums.HASH_TAG);
        keyList.add(System.currentTimeMillis() / 1000 + RateLimitConstantEnums.HASH_TAG);
        String toString = redisTemplate.execute(redisScript, keyList, keyList).toString();
        if (RateLimitConstantEnums.REDIS_ERROR.equals(toString)) {
            throw new RateLimitException(message);
        }

    }
}
